#! /usr/bin/env python
#
# Implement cookbook using pyfitsio
import sys
import os

import ctypes

# Use just the C-types bindings to CFITSIO
from fitsio import cfitsio, pyfitsio

def get_string_array2(size):

    arr = (ctypes.c_char_p * size)()
    #for n in range(size):
        #temp = ctypes.create_string_buffer('\000' * cfitsio.FLEN_VALUE)
        #temp.value='0'
        #arr[n] = (ctypes.c_char * cfitsio.FLEN_VALUE)()    
    return arr

def get_string_array(size):

    arr = (ctypes.c_char_p * size)()
        
    # Need to allocate space for each arr[]. What is the best way?
    # This seems to work well. No memory leaks apparent
    # Can this be done another way?
    temps = []
    for n in range(size):
        temps.append(ctypes.create_string_buffer('\000' * cfitsio.FLEN_VALUE))
    for n in range(size):
        arr[n] = ctypes.addressof(temps[n])
        #print n, arr[n], ctypes.addressof(temps[n])

    # need to return reference to the array otherwise it goes out of scope
    # and the allocated memory is freed up. This gives a segmentation fault

    return temps, arr

def read_ttypes(ff):

    status = ctypes.c_int(0)
    numcols = int(ff.read_key('TFIELDS'))
    
    # Read the column names from the TTYPEn keywords
    nfound = ctypes.c_int()
    temps, ttype = get_string_array(numcols)

    cfitsio.fits_read_keys_str(ff.fptr, "TTYPE", 1, numcols, 
                               ttype, 
                               ctypes.byref(nfound),
                               ctypes.byref(status))
    print " Row %10s %10s %10s" % (ttype[0], ttype[1], ttype[2])

    pyfitsio.check_error(status)


def read_columns(ff):

    naxis1, numrows = ff.read_axes()

    status = ctypes.c_int(0)

    frow = ctypes.c_longlong(1)
    felem = ctypes.c_longlong(1)
    nelem = ctypes.c_longlong(numrows)
    strnull = ' '
    longnull = ctypes.c_long(0)
    floatnull = ctypes.c_float(0.0)
    anynull = ctypes.c_int()
    
    temps, name = get_string_array(numrows)

    cfitsio.fits_read_col(ff.fptr, cfitsio.TSTRING, 1, frow, felem, nelem, \
                              strnull, name, ctypes.byref(anynull), \
                              ctypes.byref(status))
    
    dia = (ctypes.c_long * numrows)()
    cfitsio.fits_read_col(ff.fptr, cfitsio.TLONG, 2, frow, felem, nelem, \
                              strnull, dia, ctypes.byref(anynull), \
                              ctypes.byref(status))
    
    den = (ctypes.c_float * numrows)()
    cfitsio.fits_read_col(ff.fptr, cfitsio.TFLOAT, 3, frow, felem, nelem, \
                              strnull, den, ctypes.byref(anynull), \
                              ctypes.byref(status))

    pyfitsio.check_error(status)

    # temps go out of scope, so problems after calling
    names = []
    for n in range(numrows): names.append(name[n])
    return names, dia, den

def main():
    
    filename = 'atestfil.fit'
    ff = pyfitsio.open(filename)

    nhdus = ff.get_num_hdus()

    for hdunum in range(1, nhdus+1):

        hdutype = ff.move_to_hdu(hdunum)
        naxes = ff.read_axes()
        numrows = naxes[1]
        print '# HDU', hdunum, hdutype, naxes

        if hdutype == cfitsio.ASCII_TBL:
            print "Reading ASCII table in HDU#", hdunum
        elif hdutype == cfitsio.BINARY_TBL:
            print "Reading binary table in HDU#", hdunum
        else:
            print "This HDU is not an ASCII or binary table. Skipping"
            continue

        read_ttypes(ff)

        name, dia, den = read_columns(ff)

        for ii in range(numrows):
            print "%5d %10s %10ld %10.2f" % (ii+1, name[ii], dia[ii], den[ii])

    ff.close()

if __name__ == '__main__':
    
    main()
        
